from __future__ import with_statement
from sets import Set
import kuai.strtool
import os.path

FORCE_FIELD_FUNCTIONS = {}
FORCE_FIELD_DEFINATIONS = {}

class ForceFieldFunction:
    def uniform(self, atoms, ks=None, indexes=None):
        raise NotImplementedError()

    def list_atoms(self, mol):
        raise NotImplementedError()

    def estimate(self, atoms):
        raise NotImplementedError()

    def rank(self):
        raise NotImplementedError()
    
    def unit(self, i):
        raise NotImplementedError()
    
    def list_keys(self, term, mol):
        result = Set()
        for i in self.list_atoms(mol):
            item = ForceFieldKey(term, [j.type for j in i])
            self.uniform(item.atoms, None, None)
            result.add(item)
        return result

    def check(self, atoms, ks, indexes):
        assert len(atoms) == self.rank()

class NonbondFunction(ForceFieldFunction):
    def uniform(self, atoms, ks, indexes):
        if len(atoms) > 1 and atoms[0] > atoms[1]:
            atoms[0], atoms[1] = atoms[1], atoms[0]
            if indexes is not None:
                indexes[0], indexes[1] = indexes[1], indexes[0]

    def list_atoms(self, mol):
        from kuai.inner import list_nonbond_atoms
        return list_nonbond_atoms(mol)
    
    def list_keys(self, term, mol, result = {}):
        atoms = [i.type for i in mol.atoms]
        result = Set()
        for i in atoms:
            item = ForceFieldKey(term, [i])
            self.uniform(item.atoms, None, None)
            result.add(item)
        return result

class AtomFunction(ForceFieldFunction): 
    def uniform(self, atoms, ks, indexes):
        pass

    def list_atoms(self, mol): 
        result = []
        for i in mol.atoms:
            result.append([i])
        return result

    def rank(self):
        return 1

class BondFunction(ForceFieldFunction):
    def uniform(self, atoms, ks, indexes):
        if atoms[0] > atoms[1]:
            atoms[0], atoms[1] = atoms[1], atoms[0]
            if indexes is not None:
                indexes[0], indexes[1] = indexes[1], indexes[0]

    def list_atoms(self, mol):
        from kuai.inner import list_atom_pair12
        return list_atom_pair12(mol)

    def rank(self):
        return 2
     
class AngleFunction(ForceFieldFunction):
    def uniform(self, atoms, ks, indexes):
        if atoms[0] > atoms[2]:
            atoms[0], atoms[2] = atoms[2], atoms[0]
            if indexes is not None:
                indexes[0], indexes[2] = indexes[2], indexes[0]

    def list_atoms(self, mol):
        from kuai.inner import list_angle_atoms
        return list_angle_atoms(mol)

    def rank(self):
        return 3
    
class TorsionFunction(ForceFieldFunction):
    def uniform(self, atoms, ks, indexes):
        if atoms[0] > atoms[3] or atoms[0] == atoms[3] and atoms[1] > atoms[2]:
            atoms[0], atoms[3] = atoms[3], atoms[0]
            atoms[1], atoms[2] = atoms[2], atoms[1]
            if indexes is not None:
                indexes[0], indexes[3] = indexes[3], indexes[0]
                indexes[1], indexes[2] = indexes[2], indexes[1]

    def list_atoms(self, mol):
        from kuai.inner import list_torsion_atoms
        return list_torsion_atoms(mol)

    def rank(self):
        return 4

class Fork0Function(ForceFieldFunction):
    def uniform(self, atoms, ks, indexes):
        if atoms[1] > atoms[2]:
            atoms[1], atoms[2] = atoms[2], atoms[1]
            if indexes is not None:
                indexes[1], indexes[2] = indexes[2], indexes[1]
        if atoms[2] > atoms[3]:
            atoms[2], atoms[3] = atoms[3], atoms[2]
            if indexes is not None:
                indexes[2], indexes[3] = indexes[3], indexes[2]
        if atoms[1] > atoms[2]:
            atoms[1], atoms[2] = atoms[2], atoms[1]
            if indexes is not None:
                indexes[1], indexes[2] = indexes[2], indexes[1]    
        
    def list_atoms(self, mol):
        from kuai.inner import list_fork0_atoms
        return list_fork0_atoms(mol)

    def rank(self):
        return 4

class Fork1Function(Fork0Function):
    def uniform(self, atoms, ks, indexes):
        atoms[0], atoms[1] = atoms[1], atoms[0]
        if indexes is not None:
            indexes[0], indexes[1] = indexes[1], indexes[0]
        Fork0Function.uniform(self, atoms, ks, indexes)
        atoms[0], atoms[1] = atoms[1], atoms[0]
        if indexes is not None:
            indexes[0], indexes[1] = indexes[1], indexes[0]
        
    def list_atoms(self, mol):
        from kuai.inner import list_fork1_atoms
        return list_fork1_atoms(mol)

    def rank(self):
        return 4

class Fork2Function(Fork0Function):
    def uniform(self, atoms, ks, indexes):
        atoms[0], atoms[2] = atoms[2], atoms[0]
        if indexes is not None:
            indexes[0], indexes[2] = indexes[2], indexes[0]
        Fork0Function.uniform(self, atoms, ks, indexes)
        atoms[0], atoms[1], atoms[2] = atoms[1], atoms[2], atoms[0]
        if indexes is not None:
            indexes[0], indexes[1], indexes[2] = indexes[1], indexes[2], indexes[0]

    def list_atoms(self, mol):
        from kuai.inner import list_fork2_atoms
        return list_fork2_atoms(mol)

    def rank(self):
        return 4

class Fork3Function(Fork0Function):
    def uniform(self, atoms, ks, indexes):
        atoms[0], atoms[3] = atoms[3], atoms[0]
        if indexes is not None:
            indexes[0], indexes[3] = indexes[3], indexes[0]
        Fork0Function.uniform(self, atoms, ks, indexes)
        atoms[0], atoms[1], atoms[2], atoms[3] = atoms[1], atoms[2], atoms[3], atoms[0] 
        if indexes is not None:
            indexes[0], indexes[1], indexes[2], indexes[3] = indexes[1], indexes[2], indexes[3], indexes[0]

    def list_atoms(self, mol):
        from kuai.inner import list_fork3_atoms
        return list_fork3_atoms(mol)

    def rank(self):
        return 4

class FunctionATYPE(AtomFunction):
    UNITS = ["e", "g/mol"]
    def unit(self, i):
        return FunctionATYPE[i]
    
class FunctionN12_6(NonbondFunction):
    UNITS = ["A", "kJ/mol"]
    def unit(self, i):
        return FunctionN12_6.UNITS[i]
    
class FunctionBHARM(BondFunction):
    UNITS = ["A", "kJ/A^2*mol"]
    def unit(self, i):
        return FunctionBHARM.UNITS[i]
    
class FunctionAHARM(AngleFunction):
    UNITS = ["rad", "kJ/rad^2*mol"]
    def unit(self, i):
        return FunctionAHARM.UNITS[i]

class FunctionTCOSS(TorsionFunction):
    def units(self, i):
        return "kJ/mol"

class FunctionTCOSP(TorsionFunction):
    def units(self, i):
        return "kJ/mol"
    
class FunctionICOSS(Fork2Function):
    def units(self, i):
        return "kJ/mol"

class FunctionIBCOS(Fork2Function):
    def units(self, i):
        return "kJ/mol"

FORCE_FIELD_FUNCTIONS["ATYPE"] = FunctionATYPE()
FORCE_FIELD_FUNCTIONS["N12_6"] = FunctionN12_6()
FORCE_FIELD_FUNCTIONS["BHARM"] = FunctionBHARM()
FORCE_FIELD_FUNCTIONS["AHARM"] = FunctionAHARM()
FORCE_FIELD_FUNCTIONS["TCOSS"] = FunctionTCOSS()
FORCE_FIELD_FUNCTIONS["TCOSP"] = FunctionTCOSP()
FORCE_FIELD_FUNCTIONS["IBCOS"] = FunctionICOSS()

FORCE_FIELD_DEFINATIONS["AMBER"] = ["N12_6", "BHARM", "AHARM", "TCOSS", "ICOSS", ]


class ForceFieldKey:
    """This class represent the key of a force field item. It included the term 
(BHARM, AHARM etc) and the atom types. """
    def __init__(self, term = 0, atoms = []):
        self.term = term
        self.atoms = atoms[:]

    def __eq__(self, vt):
        return self.term == vt.term and self.atoms == vt.atoms

    def __hash__(self):
        result = hash(self.term)
        for i in self.atoms:
            result ^= hash(i)
        return result
    
    def __repr__(self):
        return self.term + ": " + ", ".join(self.atoms) + ": "
    
    def __cmp__(self, key2):
        if self.term < key2.term:
            result = -1
        elif self.term > key2.term:
            result = 1
        else:
            atoms1 = self.atoms
            atoms2 = key2.atoms
            assert len(atoms1) == len(atoms2)
            if self.term in FORCE_FIELD_FUNCTIONS:
                func = FORCE_FIELD_FUNCTIONS[self.term]
                if isinstance(func, AngleFunction):
                    assert len(atoms1) == 3
                    assert len(atoms2) == 3
                    atoms1 = [atoms1[1], atoms1[0], atoms1[2]]
                    atoms2 = [atoms2[1], atoms2[0], atoms2[2]]
                elif isinstance(func, TorsionFunction):
                    assert len(atoms1) == 4
                    assert len(atoms2) == 4
                    atoms1 = [atoms1[1], atoms1[2], atoms1[0], atoms1[3],]
                    atoms2 = [atoms2[1], atoms2[2], atoms2[0], atoms2[3],]
                elif isinstance(func, Fork3Function):
                    assert len(atoms1) == 4
                    assert len(atoms2) == 4
                    atoms1 = [atoms1[3], atoms1[0], atoms1[1], atoms1[2],]
                    atoms2 = [atoms2[3], atoms2[0], atoms2[1], atoms2[2],]
                elif isinstance(func, Fork2Function):
                    assert len(atoms1) == 4
                    assert len(atoms2) == 4
                    atoms1 = [atoms1[2], atoms1[0], atoms1[1], atoms1[3],]
                    atoms2 = [atoms2[2], atoms2[0], atoms2[1], atoms2[3],]
                elif isinstance(func, Fork1Function):
                    assert len(atoms1) == 4
                    assert len(atoms2) == 4
                    atoms1 = [atoms1[1], atoms1[0], atoms1[2], atoms1[3],]
                    atoms2 = [atoms2[1], atoms2[0], atoms2[2], atoms2[3],]
            
            assert len(atoms1) == len(atoms2)
            result = 0
            for i in range(len(atoms1)):
                if atoms1[i] < atoms2[i]:
                    result = -1
                    break
                elif atoms1[i] > atoms2[i]:
                    result = 1
                    break
        return result
        
        
class ForceFieldItem(list):
    def __init__(self):
        list.__init__(self)
        self.tag = ""
        self.fixed = []
        self.version = self.reference = None
        
    def __str__(self):
        result = ''
        for i in range(len(self)):
            result += str(self[i])
            if i in self.fixed:
                result += '*'
            result += ', '
        result = result[:-2] + ' '
        if self.version:
            result += ":Ver" + str(self.version)
        if self.reference:
            result += ":Ref" + str(self.reference)
        if self.tag != "":
            result += self.tag
        return result
    
    @staticmethod
    def parse(line):
        tokens = [i.strip() for i in line.split(':')]
        if len(tokens) >= 3:
            try:
                pars = [j.strip() for j in tokens[2].split(",")]
                result = ForceFieldItem()
                for i in range(len(pars)):
                    if pars[i].endswith('*'):
                        result.fixed.append(i)
                        pars[i] = pars[i][:-1]
                    result.append(float(pars[i]))
                for i in tokens[3:]:
                    if i.startswith("Ref"):
                        result.reference = int(i[3:])
                    elif i.startswith("Ver"):
                        result.version = int(i[3:])
                    else:
                        result.tag += ":"+i
                return result
            except Exception, error:
                print error
        raise ValueError("Can not parse Force Field Paramter in line '%s'" % line)
    
class ForceField(dict):
    def __init__(self, type):
        dict.__init__(self)
        self.type = type
        
    def terms(self):
        """return all term of this force field"""
        from sets import Set
        return Set([i.term for i in self.keys()])
    
    def get_ATYPE(self):
        termlist = []
        for i in self.keys():
            if i.term == "ATYPE":
                termlist.append(i.atoms)
        return termlist
    
    def get_items(self, atom_list):
        result = []
        for i in self.keys():
            counter = 0
            for j in i.atoms:
                if j in atom_list:
                    counter += 1
            if counter > 0:
                result.append(i)
        return result
    
    @staticmethod
    def list_keys(mol, type):
        try:
            return FORCE_FIELD_FUNCTIONS[type].list_keys(type, mol)
        except (TypeError, KeyError):
            pass
        
        try:
            return ForceField.list_keys(FORCE_FIELD_DEFINATIONS[type])
        except (TypeError, KeyError):
            pass
        
        result = []
        for i in type:
            if i in FORCE_FIELD_FUNCTIONS:
                result += ForceField.list_keys(mol, i)
        return result
        
    def __repr__(self):
        from kuai.dff.dffppf import writer
        from cStringIO import StringIO
        file = StringIO()
        writer.write(file, self)
        result = file.getvalue()
        return result
    

class EquivalenceTableItem(list):
    
    COLUMNS = 10
    
    def __init__(self):
        list.__init__(self)
        self.tag = ""
        self.version = self.reference = None
        
    def __repr__(self):
        result = ''
        result += " ".join(self) + " "
        if self.version:
            result += ":Ver" + str(self.version)
        if self.reference:
            result += ":Ref" + str(self.reference)
        result += self.tag
        return result
    
    @staticmethod
    def parse(line):
        tokens = line.split(':')
        if len(tokens) >= 2:
            pars = tokens[1].split()
            if len(pars) >= EquivalenceTableItem.COLUMNS:
                try:
                    result = EquivalenceTableItem()
                    result += pars
                    for i in tokens[2:]:
                        if i.startswith("Ref"):
                            result.reference = int(i[3:])
                        elif i.startswith("Ver"):
                            result.version = int(i[3:])
                        else:
                            self.tag += ":"+i
                    
                    return result
                except Exception:
                    pass
        raise ValueError("Can not parse Equivalence Table Paramter in line '%s'" % line)
    
class EquivalenceTable(dict):
    
    EQT_FILE_HEADER = """#AAT             : NB    ATC    BINC    Bond    A/C    A/S    T/C    T/S    O/C    O/S
"""
    
    def __init__(self):
        return dict.__init__(self)
    
    def get_equal_type(self, term, i, type):
        from kuai.charges import FunctionATC, FunctionBINC
        if type in self:
            try:
                func = FORCE_FIELD_FUNCTIONS[term]
                if isinstance(func, NonbondFunction):
                    return self[type][0]
                elif isinstance(func, FunctionATC):
                    return self[type][1]
                elif isinstance(func, FunctionBINC):
                    return self[type][2]
                elif isinstance(func, BondFunction):
                    return self[type][3]
                elif isinstance(func, AngleFunction):
                    if i == 1:
                        return self[type][4]
                    else:
                        return self[type][5]
                elif isinstance(func, TorsionFunction):
                    if i == 2 or i == 3:
                        return self[type][6]
                    else:
                        return self[type][7]
                elif isinstance(func, Fork3Function):
                    if i == 3:
                        return self[type][8]
                    else:
                        return self[type][9]
                elif isinstance(func, Fork2Function):
                    if i == 2:
                        return self[type][8]
                    else:
                        return self[type][9]
                elif isinstance(func, Fork1Function):
                    if i == 1:
                        return self[type][8]
                    else:
                        return self[type][9]
                elif isinstance(func, Fork0Function):
                    if i == 0:
                        return self[type][8]
                    else:
                        return self[type][9]
            except KeyError:
                pass
            
        return type
                
    def get_equal_key(self, item):
        atoms = [self.get_equal_type(item.term, i, item.atoms[i]) for i in range(len(item.atoms))]
        return ForceFieldKey(item.term, atoms)
    
    def get_equal_force_field(self, ff):
        result = ForceField(ff.type)
        for k, v in ff.iteritems():
            result[self.get_equal_key(k)] = v
        return result
    
    def __repr__(self):
        result = EquivalenceTable.EQT_FILE_HEADER
        keys = [i for i in self.iterkeys()]
        keys.sort()
        for k in keys:
            result += k + ": " + str(self[k]) + "\n"
        return result 

    @staticmethod
    def read(file):
        result = EquivalenceTable()
        for i in file:
            i = i.strip()
            if not i.startswith('#'):
                try:
                    item = EquivalenceTableItem.parse(i)
                    key = i.split(':')[0].strip()
                    result[key] = item
                except Exception:
                    pass                
        return result
    
    @staticmethod
    def get_eqt_filename(file):
        folder, filename = os.path.split(file)
        ext = kuai.strtool.extname(filename)
        if ext == '':
            filename += '.eqt'
        else:
            filename = filename[:-len(ext)]+"eqt"
        result = os.path.join(folder, filename)
        return result
    
def checkout(ff, eqt, mol):
    """This function return the required force field and equivalence table for mol"""
    keysets = []
    for i in FORCE_FIELD_DEFINATIONS[ff.type]:
        try:
            func = FORCE_FIELD_FUNCTIONS[i]
            keysets += func.list_keys(mol, i)
        except KeyError:
            pass
    resultEQT = EquivalenceTable()
    for i in mol.atoms:
        if i.type in eqt:
            resultEQT[i.type] = eqt[i.type]
    resultFF = ForceField(ff.type)
    for i in keysets:
        newkey = resultEQT.get_equal_key(i)
        try:
            resultFF[newkey] = ff[newkey]
        except KeyError:
            pass
    return (resultFF, resultEQT)

def check_coverage(ff, eqt, mol):
    """This function return the required """
    missings = []
    for i in FORCE_FIELD_DEFINATIONS[ff.type][:4]:
        try:
            func = FORCE_FIELD_FUNCTIONS[i]
            for j in func.list_keys(mol, i):
                newkey = eqt.equal_item(j)
                if newkey not in ff:
                    missings.append(newkey)
        except KeyError:
            pass
    if len(missings) > 0:
        error = KeyError("There are missing items in force field!")
        error.missings = missings
        raise error

